package com.lzwork.gpt.utility.common;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONWriter.Feature;
import com.lzwork.utility.file.LineContent;

/**
 * @author LiZheng
 * @date 2023年9月25日 下午2:23:48
 */

public class TrainMaker
{
	public TrainMaker()
	{
		this("instruction", "input", "output");
	}

	public TrainMaker(String keyInst, String keyInput, String keyOutput)
	{
		this.keyInst = keyInst;
		this.keyInput = keyInput;
		this.keyOutput = keyOutput;
	}

	protected String keyInst;
	protected String keyInput;
	protected String keyOutput;

	List<Unit> dataList = new ArrayList<>();

	class Unit
	{
		Unit(String inst, String input, String output)
		{
			this.inst = inst;
			this.input = input;
			this.output = output;
		}

		String inst;
		String input;
		String output;

		JSONObject toJObj()
		{
			JSONObject obj = new JSONObject(true);
			obj.put(keyInst, inst);
			obj.put(keyInput, input);
			obj.put(keyOutput, output);
			return obj;
		}
	}

	public void append(String inst, String input, String output)
	{
		dataList.add(new Unit(inst, input, output));
	}

	public boolean writeFile(File output, boolean asLines)
	{
		List<JSONObject> objList = new ArrayList<>();

		for (Unit unit : dataList)
		{
			objList.add(unit.toJObj());
		}

		List<String> toWrite = new ArrayList<>();
		if (asLines)
		{
			for (JSONObject o : objList)
			{
				toWrite.add(o.toJSONString());
			}
		} else
		{
			JSONArray ja = new JSONArray(objList);

			String out = ja.toJSONString(Feature.PrettyFormat);
			toWrite.add(out);
		}

		return LineContent.saveFile(output, toWrite);
	}

	public int size()
	{
		return dataList.size();
	}
}
